import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Joy
from unitree_go.msg import WirelessController

class UJoyConverter(Node):

    def __init__(self):
        super().__init__('ujoy_converter')
        self.subscription = self.create_subscription(WirelessController, '/wirelesscontroller', self.listener_callback, 10)
        self.publisher_ = self.create_publisher(Joy, '/joy', 10)
        self.joy = Joy()
        self.key_last = 0
    
    def listener_callback(self, wlc_msg: WirelessController):
        # self.get_logger().info(f'{wlc_msg.lx:.1f}\n{wlc_msg.ly:.1f}\n{wlc_msg.rx:.1f}\n{wlc_msg.ry:.1f}\n{wlc_msg.keys:.1f}\n---')
        key_press = self.anti_continuous_pressing(wlc_msg.keys)  # 防连按
        # self.get_logger().info(f'key press: {key_press}')

        self.msg = Joy()
        # Joy()默认是空数组，要初始化
        self.msg.axes = [0.0] * 4
        self.msg.buttons = [0] * 12
        # 模拟 Xinput 输出
        self.msg.axes[1] = wlc_msg.ly  # to x
        self.msg.axes[0] = wlc_msg.lx  # to y
        self.msg.axes[3] = wlc_msg.rx  # to yaw
        # 防连按处理
        if key_press == 2:
            self.msg.buttons[4] = 1  # LB
        elif key_press == 1:
            self.msg.buttons[5] = 1  # RB
        elif key_press == 3:
            self.msg.buttons[4] = 1  # LB
            self.msg.buttons[5] = 1  # RB
        else:
            self.msg.buttons[4] = self.msg.buttons[5] = 0
        
        self.publisher_.publish(self.msg)
    
    def anti_continuous_pressing(self, key):
        if self.key_last != key:
            key_press = key
            self.key_last = key
        else:
            key_press = 0
        return key_press

def main():
    rclpy.init()

    ujoy_converter = UJoyConverter()
    ujoy_converter.get_logger().info('Ujoy converter is running!')
    rclpy.spin(ujoy_converter)

    rclpy.shutdown()

if __name__ == '__main__':
    main()
